package org.groupg.practice;

import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.io.jdbc.JdbcIO;
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.join.CoGbkResult;
import org.apache.beam.sdk.transforms.join.CoGroupByKey;
import org.apache.beam.sdk.transforms.join.KeyedPCollectionTuple;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TypeDescriptors;
import org.apache.commons.dbcp2.BasicDataSource;

public class BeamJdbcJoinExample {

    public static void main(String[] args) {
        Pipeline pipeline = Pipeline.create();


        BasicDataSource dataSource = new BasicDataSource();
        dataSource.setDriverClassName("com.mysql.cj.jdbc.Driver");
        dataSource.setUrl("jdbc:mysql://localhost:3306/mydatabase");
        dataSource.setUsername("username");
        dataSource.setPassword("password");
        dataSource.setMinIdle(5);
        dataSource.setMaxIdle(10);
        dataSource.setMaxTotal(20);
        dataSource.setMaxWaitMillis(10000);

        // 使用 JdbcIO 读取两个数据源
        PCollection<KV<String, String>> dataset1 = pipeline.apply(
                JdbcIO.<KV<String, String>>read()
                        .withDataSourceConfiguration(JdbcIO.DataSourceConfiguration.create(dataSource))
                        .withQuery("SELECT id, name FROM table1")
                        .withCoder(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()))
                        .withRowMapper(row -> KV.of(row.getString(0), row.getString(1))));

        PCollection<KV<String, String>> dataset2 = pipeline.apply(
                JdbcIO.<KV<String, String>>read()
                        .withDataSourceConfiguration(JdbcIO.DataSourceConfiguration.create(dataSource))
                        .withQuery("SELECT id, name FROM table2")
                        .withCoder(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()))
                        .withRowMapper(row -> KV.of(row.getString(0), row.getString(1))));

        // 定义用于 join 的标签
        final TupleTag<String> tag1 = new TupleTag<>();
        final TupleTag<String> tag2 = new TupleTag<>();

        // 使用 CoGroupByKey 进行 join
        PCollection<KV<String, CoGbkResult>> joined = KeyedPCollectionTuple
                .of(tag1, dataset1)
                .and(tag2, dataset2)
                .apply(CoGroupByKey.<String>create());

        // 处理 join 结果
        PCollection<String> result = joined.apply(
                MapElements.into(TypeDescriptors.strings())
                        .via((KV<String, CoGbkResult> element) -> {
                            Iterable<String> values1 = element.getValue().getAll(tag1);
                            Iterable<String> values2 = element.getValue().getAll(tag2);

                            // 实现 join 逻辑，例如内连接
                            StringBuilder sb = new StringBuilder();
                            for (String v1 : values1) {
                                for (String v2 : values2) {
                                    sb.append(v1).append(",").append(v2).append("\n");
                                }
                            }
                            return sb.toString();
                        }));

        // 输出结果（可以选择输出到文件或其他存储）
        // result.apply(...);



        pipeline.run().waitUntilFinish();
    }
}